import cv2
import torch
import numpy as np
from collections import deque
from model import ResLstmNet
from DataSet import data_transform
from PIL import Image
from dection import check_image
from ultralytics import YOLO

sequence_len = 15
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = ResLstmNet().to(device)
model.load_state_dict(torch.load("./runs/gunshot79.pth", map_location=device))
# model.load_state_dict(torch.load("model/913_764/gunshot27.pth", map_location=device))
yolo_model = YOLO(r'H:\shooting\model\gun_check_V2.pt')
model.eval()


# cap = cv2.VideoCapture(r'E:\Users\kx15\Desktop\practise\shooting\left_2.mp4')
# cap = cv2.VideoCapture(r'E:\Users\kx15\Desktop\practise\shooting\left_3.mp4')
# cap = cv2.VideoCapture(r'E:\Users\kx15\Desktop\practise\shooting\left_4.mp4')
cap = cv2.VideoCapture(r'E:\Users\kx15\Desktop\practise\shooting\left_5.mp4')

frames = deque(maxlen=sequence_len)

gunshot_detected = False
behavior_count = 0

with torch.no_grad():
    while True:
        ret, frame = cap.read()
        if not ret:
            break

        show_frame = frame.copy()
        frame, _ = check_image(frame, yolo_model)
        if frame is None:
            continue
        # 预处理单帧
        img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        img = Image.fromarray(img)
        img = data_transform(img)
        frames.append(img)

        if len(frames) == sequence_len:
            seq = torch.stack(list(frames))  # [T, C, H, W]
            seq = seq.unsqueeze(0).to(device)  # [1, T, C, H, W]

            output = model(seq)  # 推理
            # pred = (output > 0.5).float().item()  # 二分类
            _, pred_class = torch.max(output, 1)  # 多分类
            pred = pred_class.item()  # 多分类

            # label = "Gunshot" if pred == 1 else "Normal"  # 二分类

            class_map_inv = {0: "unfire", 1: "fire", 2: "stop"}  # 多分类
            label = class_map_inv[pred]  # 多分类
# -------------------------------------------------------------
            # %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
            # # 二分类
            # if label == "Normal" and not gunshot_detected:
            #     gunshot_detected = True
            #
            # if label == "Gunshot" and gunshot_detected:
            #     gunshot_detected = False
            #     behavior_count += 1  # 行为完成，计数器加1
            #     print(f"Behavior completed! Total behaviors: {behavior_count}")
            # %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
            # 多分类
            if label == "unfire" and not gunshot_detected:
                gunshot_detected = True

            if label == "fire" and gunshot_detected:
                gunshot_detected = False
                behavior_count += 1  # 行为完成，计数器加1
                print(f"Behavior completed! Total behaviors: {behavior_count}")
            if label == "stop" and behavior_count != 0 and behavior_count >= 11:
                print("检测到结束标识！")
                break
# -------------------------------------------------------------
            show_frame = cv2.resize(show_frame, (0, 0), fx=0.3, fy=0.3)

            cv2.putText(show_frame, f"Prediction: {label}", (30, 50),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

            cv2.putText(show_frame, f"Behaviors: {behavior_count}", (30, 100),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

        cv2.imshow("Real-time Detection", show_frame)

        if cv2.waitKey(1) & 0xFF == 27:
            break

cap.release()
cv2.destroyAllWindows()
